package com.rem.flink.flink8State;

import com.rem.flink.flink2Source.ClickSource;
import com.rem.flink.flink2Source.Event;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.AggregatingState;
import org.apache.flink.api.common.state.AggregatingStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;

/**
 * 按键分区状态 AggregatingState
 * 统计每个用户点击频率，到达五次就输出统计结果
 *
 * @author Rem
 * @date 2022-10-25
 */

public class AggregatingStateTest {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(WatermarkStrategy.<Event>forMonotonousTimestamps()
                        .withTimestampAssigner((element, recordTimestamp) -> element.getTimestamp()));

        stream.printToErr("time");


        stream.keyBy(Event::getUser)
                .flatMap(new RichFlatMapFunction<Event, String>() {
                    // 定义一个值状态，用来保存当前用户访问频次
                    ValueState<Long> countState;
                    // 定义聚合状态，用来计算平均时间戳
                    AggregatingState<Event, Long> avgTsAggState;


                    @Override
                    public void open(Configuration parameters) {
                        countState = getRuntimeContext().getState(new ValueStateDescriptor<>("count", Long.class));

                        avgTsAggState = getRuntimeContext().getAggregatingState(new AggregatingStateDescriptor<>("avg-ts",
                                        /**
                                         * 入参 Event
                                         * 中间状态 Tuple2<时间戳，访问次数>
                                         * 返回：频率
                                         */
                                        new AggregateFunction<Event, Tuple2<Long, Long>, Long>() {

                                            @Override
                                            public Tuple2<Long, Long> createAccumulator() {
                                                return Tuple2.of(0L, 0L);
                                            }

                                            @Override
                                            public Tuple2<Long, Long> add(Event value, Tuple2<Long, Long> accumulator) {
                                                return Tuple2.of(accumulator.f0 + value.getTimestamp(), accumulator.f1 + 1);
                                            }

                                            @Override
                                            public Long getResult(Tuple2<Long, Long> accumulator) {
                                                return accumulator.f0 / accumulator.f1;
                                            }

                                            @Override
                                            public Tuple2<Long, Long> merge(Tuple2<Long, Long> a, Tuple2<Long, Long> b) {
                                                return null;
                                            }
                                        },
                                        Types.TUPLE(Types.LONG, Types.LONG)
                                )
                        );
                    }

                    @Override
                    public void flatMap(Event value, Collector<String> out) throws Exception {
                        Long count = countState.value();
                        if (count == null) {
                            count = 1L;
                        } else {
                            count++;
                        }
                        countState.update(count);
                        avgTsAggState.add(value);

                        // 达到5次就输出结果，并清空状态
                        if (count == 5) {
                            out.collect(value.getUser() + " 平均时间戳：" + new Timestamp(avgTsAggState.get()));
                            countState.clear();
                        }
                    }
                })
                .print();


        env.execute();
    }


}
